import torch
import argparse


def get_args_parser():
    parser = argparse.ArgumentParser(description='V-coco eval')
    # general
    parser.add_argument('--ckpt',
                        help='experiment configure file name',
                        required=True,
                        type=str)

    args = parser.parse_args()

    return args


def extract_parameter(ckpt_path):
    name = ckpt_path.split('/')[-1]
    ckpt_model = torch.load(ckpt_path, map_location=torch.device('cpu'))
    hoi_pose_net = ckpt_model['model_state_dict']
    pose_net = hoi_pose_net.copy()
    print(type(hoi_pose_net))
    # 'self.pose_net'
    for key, val in hoi_pose_net.items():
        if 'pose_net' in key:
            new_key = key.split('pose_net.')[-1]
            print(f'{key} -> {new_key}')
            pose_net[new_key] = val
        pose_net.pop(key)
    torch.save(pose_net, f'pose_net_from_{name}')


if __name__ == "__main__":
    args = get_args_parser()
    extract_parameter(args.ckpt)
